﻿//
// Created by QU on 24-9-28.
//

// TODO: When sending split parts to server, server will reply the full merged node.
//       But why cause server crash?
#include "Stress_client.h"
#include "node/Node.h"

#include <iostream>
#include <thread>
#include <boost/endian/conversion.hpp>

StressClient::StressClient(const std::string& host, const std::string& port, const Config& config)
    : m_host(host), m_port(port), m_config(config)
{
}

void StressClient::startLoad()
{
    std::vector<std::future<void>> futures;
    std::atomic<bool> should_stop { false };
    auto start_time = std::chrono::steady_clock::now();

    for (int i = 0; i < m_config.thread_count; ++i)
    {
        futures.emplace_back(std::async(std::launch::async, [this, &should_stop, start_time]()
        {
            runClientThread(should_stop, start_time);
        }));
    }

    for (auto& future : futures)
    {
        future.wait();
    }

    std::cout << "All threads completed. Total threads: " << m_config.thread_count << std::endl;
    std::cout << "Total sent messages: " << m_total_sent_messages << std::endl;
    std::cout << "Total received messages: " << m_total_received_messages << std::endl;
}

void StressClient::runClientThread(std::atomic<bool>& should_stop, std::chrono::steady_clock::time_point start_time)
{
    try
    {
        boost::asio::io_context io_context;
        boost::asio::ip::tcp::socket socket(io_context);
        connect(socket);

        std::random_device rd;
        std::default_random_engine gen(rd());
        std::uniform_int_distribution<> dis(0, 3);

        for (int i = 0; i < m_config.messages_per_thread && !should_stop; ++i)
        {
            sendRandomMessage(socket, dis(gen));
            receiveEcho(socket);
            delayBetweenMessages(gen);

            if (std::chrono::steady_clock::now() - start_time >= m_config.total_test_duration)
            {
                should_stop = true;
            }
        }
    }
    catch (const std::exception& e)
    {
        std::cerr << "Thread " << std::this_thread::get_id() << " exception: " << e.what() << std::endl;
    }
}

void StressClient::connect(boost::asio::ip::tcp::socket& socket) const
{
    boost::asio::ip::tcp::resolver resolver(socket.get_executor());
    auto endpoints = resolver.resolve(m_host, m_port);
    boost::asio::connect(socket, endpoints);
}

void StressClient::sendRandomMessage(boost::asio::ip::tcp::socket& socket, int message_type)
{
    switch (message_type)
    {
    case 0: sendSmallMessage(socket); break;
    case 1: sendLargeMessage(socket); break;
    case 2: sendMultipleMessages(socket); break;
    case 3: sendSplitMessage(socket); break;
    }
}

void StressClient::sendSmallMessage(boost::asio::ip::tcp::socket& socket)
{
    transmitMessage(socket, randomSize(m_config.small_message_min, m_config.small_message_max), "small");
}

void StressClient::sendLargeMessage(boost::asio::ip::tcp::socket& socket)
{
    transmitMessage(socket, randomSize(m_config.large_message_min, m_config.large_message_max), "large");
}

void StressClient::sendMultipleMessages(boost::asio::ip::tcp::socket& socket)
{
    for (int i = 0; i < m_config.multiple_message_count; ++i)
    {
        transmitMessage(socket, randomSize(m_config.multiple_message_min, m_config.multiple_message_max),
                        "multiple part " + std::to_string(i + 1));
    }
}

void StressClient::sendSplitMessage(boost::asio::ip::tcp::socket& socket)
{
    auto data = constructMessage(randomSize(m_config.split_message_min, m_config.split_message_max));
    size_t split_point = data.size() / 2;

    boost::asio::write(socket, boost::asio::buffer(data.data(), split_point));
    std::cout << "Thread " << std::this_thread::get_id() << " sent split message part 1: " << split_point << " bytes"
              << std::endl;

    std::this_thread::sleep_for(std::chrono::milliseconds(10));

    boost::asio::write(socket, boost::asio::buffer(data.data() + split_point, data.size() - split_point));
    std::cout << "Thread " << std::this_thread::get_id() << " sent split message part 2: " << (data.size() - split_point)
              << " bytes" << std::endl;
}

std::vector<std::byte> StressClient::constructMessage(size_t payload_size)
{
    Node node(payload_size, 0x01, static_cast<uint8_t>(MSGID::MSG_HELLO_WORLD));
    std::vector<std::byte> payload(payload_size);
    std::generate(payload.begin(), payload.end(), [&]() { return static_cast<std::byte>(m_gen()); });
    node.setPayload(payload.data(), payload_size);

    auto serialized = node.serialize();
    std::vector<std::byte> result(serialized.buffer.get(), serialized.buffer.get() + serialized.length);
    return result;
}

void StressClient::receiveEcho(boost::asio::ip::tcp::socket& socket)
{
    std::vector<char> header_buffer(Node::MAX_HEADER_SIZE);
    boost::system::error_code ec;

    size_t bytes_transferred = boost::asio::read(socket,
        boost::asio::buffer(header_buffer, Node::MAX_HEADER_SIZE), ec);

    if (ec || bytes_transferred < Node::MAX_HEADER_SIZE)
    {
        std::cerr << "Receive header error: " << ec.message() << std::endl;
        return;
    }

    const uint16_t payload_size = boost::endian::big_to_native(
        *reinterpret_cast<uint16_t*>(&header_buffer[2]));

    std::vector<std::byte> payload_buffer(payload_size);
    bytes_transferred = boost::asio::read(socket,
        boost::asio::buffer(payload_buffer, payload_size), ec);

    if (ec || bytes_transferred < payload_size)
    {
        std::cerr << "Receive payload error: " << ec.message() << std::endl;
        return;
    }

    m_total_received_messages++;
    std::cout << "Thread " << std::this_thread::get_id() << " received echo: " << payload_size << " bytes" << std::endl;
}

void StressClient::transmitMessage(boost::asio::ip::tcp::socket& socket, size_t payload_size, const std::string& message_type)
{
    auto data = constructMessage(payload_size);
    boost::asio::write(socket, boost::asio::buffer(data));
    m_total_sent_messages++;
    std::cout << "Thread " << std::this_thread::get_id() << " sent " << message_type << " message: " << data.size() << " bytes"
              << std::endl;
}

void StressClient::delayBetweenMessages(std::default_random_engine& gen)
{
    std::this_thread::sleep_for(std::chrono::milliseconds(
        std::uniform_int_distribution<>(m_config.delay_between_messages_min,
                                        m_config.delay_between_messages_max)(gen)));
}

size_t StressClient::randomSize(int min, int max)
{
    return std::uniform_int_distribution<>(min, max)(m_gen);
}

int main(int argc, char* argv[])
{
    try
    {
        StressClient::Config config;
        config.thread_count = 4;
        config.messages_per_thread = 1000;
        config.total_test_duration = std::chrono::minutes(1);

        StressClient client("127.0.0.1", "10120", config);
        client.startLoad();
    }
    catch (const std::exception& e)
    {
        std::cerr << "Exception: " << e.what() << std::endl;
    }

    return 0;
}